#include "common.h"
#include "server.h"

#include <map>

constexpr int BUFFER_SIZE = 4096;
constexpr int MAX_CLIENTS = 10;

bool
TcpServer::init()
{
    if (_socket != -1) {
        std::cout << "Socket already initialized" << std::endl;
        return false;
    }

    init_sockets();

    _socket = socket(AF_INET, SOCK_STREAM, 0);
    if (_socket == -1) {
        std::cerr << "Error creating socket" << std::endl;
        return false;
    }

    // Allow socket reuse address (optional but good for testing)
    int optval = 1;
    setsockopt(_socket, SOL_SOCKET, SO_REUSEADDR, (const char*)&optval, sizeof(optval));

    sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(_port);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (-1 == bind(_socket, (sockaddr*)&addr, sizeof(addr))) {
        std::cerr << "Error binding socket" << std::endl;
        close_socket(_socket);
        _socket = -1;
        return false;
    }

    if (-1 == listen(_socket, MAX_CLIENTS)) {
        std::cerr << "Error listening on socket" << std::endl;
        close_socket(_socket);
        _socket = -1;
        return false;
    }
    std::cout << "Server listening on port " << _port << "..." << std::endl;
    return true;
}

void 
TcpServer::start()
{
    _running = true;
    // Use a map to store the receive buffer for each connected client
    std::map<int, std::vector<char>> clientBuffers;
    while (_running) {
        sockaddr_in client_addr;
        socklen_t client_addr_len = sizeof(client_addr);
        int client_socket = accept(_socket, (sockaddr*)&client_addr, &client_addr_len);
        if (client_socket == -1) {
            std::cerr << "Error accepting connection" << std::endl;
            continue;
        }

        char client_ip[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
        std::cout << "Accepted connection from " << client_ip << ":" 
                  << ntohs(client_addr.sin_port) 
                  << " (socket " << client_socket << ")" << std::endl;

        clientBuffers[client_socket].reserve(BUFFER_SIZE * 2);
        char recvBuffer[BUFFER_SIZE];
        ssize_t bytesReceived;

        while ((bytesReceived = recv(client_socket, recvBuffer, BUFFER_SIZE, 0)) > 0) {
            clientBuffers[client_socket].insert(clientBuffers[client_socket].end(), recvBuffer, recvBuffer + bytesReceived);
            // Process data from the buffer
            processClientData(client_socket, clientBuffers[client_socket]);
        }

        if (bytesReceived == 0) 
            std::cout << "Client " << client_socket << " disconnected" << std::endl;
        else 
            std::cerr << "Error receiving data from client" <<  client_socket << ": " << strerror(errno) <<  std::endl;
        
        close_socket(client_socket);
        std::vector<char>().swap(clientBuffers[client_socket]);
    }

    close_socket(_socket);
    cleanup_sockets();
    _socket = -1;
    _running = false;
}

void 
TcpServer::stop()
{
    _running = false;
}

void 
TcpServer::processClientData(int client_socket, std::vector<char>& buffer)
{
    while (true) {
        // 1. Check if buffer contains enough data for the length prefix
        if (buffer.size() < LENGTH_PREFIX_SIZE) {
            // Not enough data for the length prefix, need to receive more
            std::cout << "Buffer too small for length prefix. Current size: " << buffer.size() << std::endl;
            break;
        }

        // 2. Read the length prefix
        uint32_t length = 0;
        memcpy(&length, buffer.data(), LENGTH_PREFIX_SIZE);
        length = ntohl(length);

        std::cout << "Potential message length: " << length << " (from buffer size " << buffer.size() << ")" << std::endl;

        // Basic validation of message length (prevent unreasonable values)
        if (length > buffer.capacity() * 2) {
            // Arbitrary limit, e.g., max possible message size
            std::cerr << "Message length too large: " << length << std::endl;
            std::vector<char>().swap(buffer);

            break;
        }

        // 3. Check if buffer contains enough data for the message
        if (buffer.size() < length + LENGTH_PREFIX_SIZE) {
            // Not enough data for the message, need to receive more
            std::cout << "Buffer too small for message. Current size: " << buffer.size() << std::endl;
            break;
        }

        // 4. Process the message
        std::cout << "Received message: " << std::string(buffer.data() + LENGTH_PREFIX_SIZE, length) << std::endl;

        // 5. Remove the processed message from the buffer
        buffer.erase(buffer.begin(), buffer.begin() + length + LENGTH_PREFIX_SIZE);
    }
}
